import os
import regex as re
from typing import BinaryIO, Counter, List, Iterator
from multiprocessing import Pool, cpu_count


class PreTokenizer:
    def __init__(self, special_tokens: List[str] = []):
        self.special_tokens = sorted(special_tokens, key=len, reverse=True)

        self.special_tokens_pattern = (
            "|".join(re.escape(token) for token in self.special_tokens)
            if self.special_tokens
            else r"(?!)"
        )  # 根据 special_token 进行分段的正则表达式，。没有特殊token时，使用永远不分割的表达式

        self.pattern = re.compile(
            r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
        )  # pretokenize 时使用的正则表达式

    def find_chunk_boundaries(
        self, file: BinaryIO, desired_num_chunks: int, split_special_token: bytes
    ) -> list[int]:
        """
        Chunk the file into parts that can be counted independently.
        May return fewer chunks if the boundaries end up overlapping.
        """
        assert isinstance(
            split_special_token, bytes
        ), "Must represent special token as a bytestring"

        # Get total file size in bytes
        file.seek(0, os.SEEK_END)
        file_size = file.tell()
        file.seek(0)

        chunk_size = file_size // desired_num_chunks

        # Initial guesses for chunk boundary locations, uniformly spaced
        # Chunks start on previous index, don't include last index
        chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
        chunk_boundaries[-1] = file_size

        mini_chunk_size = 4096  # Read ahead by 4k bytes at a time

        for bi in range(1, len(chunk_boundaries) - 1):
            initial_position = chunk_boundaries[bi]
            file.seek(initial_position)  # Start at boundary guess
            while True:
                mini_chunk = file.read(mini_chunk_size)  # Read a mini chunk

                # If EOF, this boundary should be at the end of the file
                if mini_chunk == b"":
                    chunk_boundaries[bi] = file_size
                    break

                # Find the special token in the mini chunk
                found_at = mini_chunk.find(split_special_token)
                if found_at != -1:
                    chunk_boundaries[bi] = initial_position + found_at
                    break
                initial_position += mini_chunk_size

        # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
        return sorted(set(chunk_boundaries))

    def count_pretoken_frequency(
        self, input_path: str | os.PathLike, parallel_optimize: bool = False
    ) -> Counter[bytes]:
        """统计指定文件中每个 pre_token 的出现频次"""

        if parallel_optimize:  # 启用并行优化，开启多进程处理
            subproc_cnt = cpu_count() * 2
            chunk_per_subproc = 5
        else:  # 默认单进程处理
            subproc_cnt = 1
            chunk_per_subproc = 100

        with open(input_path, "rb") as f:
            # 获取分块边界
            boundaries = self.find_chunk_boundaries(
                f, subproc_cnt * chunk_per_subproc, "<|endoftext|>".encode("utf-8")
            )

        # 开启子进程进行 pretokenize 并统计频次
        with Pool() as pool:
            results = pool.starmap(
                self._count_pretoken_frequency_by_chunks,
                [
                    (
                        input_path,
                        boundaries[
                            i * chunk_per_subproc : (i + 1) * chunk_per_subproc + 1
                        ],
                    )
                    for i in range(subproc_cnt)
                ],
            )

        # 合并所有子进程的结果
        pretoken_counter: Counter[str] = Counter()
        for res in results:
            pretoken_counter.update(res)

        # 统一做 str -> bytes 转换
        pretoken_bytes_counter: Counter[bytes] = Counter()
        for word, freq in pretoken_counter.items():
            pretoken_bytes_counter[word.encode("utf-8")] = freq

        return pretoken_bytes_counter

    def _count_pretoken_frequency_by_chunks(
        self, input_path: str | os.PathLike, boundaries: List[int]
    ):
        """对文件的指定分块统计每个 pretoken 的出现频次"""
        pt_counter: Counter[str] = Counter()

        with open(input_path, "rb") as f:
            for start, end in zip(boundaries[:-1], boundaries[1:]):
                f.seek(start)
                chunk = f.read(end - start).decode("utf-8", errors="ignore")

                # 这里选择 iter 的方式是为了节省内存，不用保存完整的 pretoken 列表
                # 选择以 str 的方式返回，是为了提高运行速度（str 相比 bytes 的哈希效率更高）
                for pt_list in self.pretokenize_iter(chunk, in_bytes=False):
                    pt_counter.update(pt_list)

        return pt_counter

    def pretokenize(
        self,
        text: str,
        preserve_special_tokens: bool = False,
    ) -> List[bytes]:
        """
        对指定的文本进行预分词，返回预分词后的字节列表，可选择是否保留 special_tokens
        """

        # 判断是否保留 special_tokens
        split_pattern = (
            f"({self.special_tokens_pattern})"
            if preserve_special_tokens
            else self.special_tokens_pattern
        )
        parts = re.split(split_pattern, text)

        res = []

        for part in parts:
            if part in self.special_tokens:
                res.append(part.encode("utf-8"))
            else:
                # 按照 self.pattern 进行正则匹配
                pretokens = [
                    token.group().encode("utf-8")
                    for token in self.pattern.finditer(part)
                ]
                res.extend(pretokens)

        return res

    def pretokenize_iter(
        self,
        text: str,
        preserve_special_tokens: bool = False,
        in_bytes: bool = True,
    ) -> Iterator[List]:
        """对指定的文本进行预分词，返回一个可迭代的预分词列表，可选择是否保留 special_tokens，结果是否以字节的形式返回"""

        # 判断是否保留 special_tokens
        split_pattern = (
            f"({self.special_tokens_pattern})"
            if preserve_special_tokens
            else self.special_tokens_pattern
        )
        parts = re.split(split_pattern, text)

        for part in parts:
            res = []
            if part in self.special_tokens:
                res = [part.encode("utf-8") if in_bytes else part]
            else:
                res = [
                    token.group().encode("utf-8") if in_bytes else token.group()
                    for token in self.pattern.finditer(part)
                ]
            yield res


if __name__ == "__main__":
    pretokenizer = PreTokenizer(["<|endoftext|>"])

    # text = "hello world <|endoftext|> my friend"
    # res = pretokenizer.pretokenize(text)
    # print(res)
